#!/usr/bin/env python
# Copyright (c) 2018-2020, NVIDIA CORPORATION. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#  * Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#  * Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#  * Neither the name of NVIDIA CORPORATION nor the names of its
#    contributors may be used to endorse or promote products derived
#    from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
# PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.


import argparse
import numpy as np
import os
from builtins import range
#from tensorrtserver.api import *
import cv2
#import tritongrpcclient
import tritonclient.http as httpclient
from tritonclient.utils import InferenceServerException
from nms.nms_py import cpu_nms
import face_align





def get_box(triton_client,image_url):
    model_name = "retinaface-50"
    model_version = 1
    batch_size = 1
    #infer_ctx = InferContext('localhost:28001', protocol, model_name, model_version,http_headers=None, verbose=False)
    img_data = cv2.imread(image_url)
    img_data = np.array(img_data)
    print(img_data.shape)
    img_data = np.transpose(img_data, [2, 0, 1])
    img_data = np.expand_dims(img_data, 0).astype('float32')
    inputs =[]
    inputs.append(httpclient.InferInput('data', [1, 3, 352, 640], "FP32"))
    outputs = []
    inputs[0].set_data_from_numpy(img_data)


    outputs.append(httpclient.InferRequestedOutput('prob'))
    results = triton_client.infer(
        model_name=model_name, inputs=inputs,  outputs=outputs, headers={"test": "1"}
    )
    boxes = results.as_numpy("prob")[0]
    print(boxes.shape)

    return boxes.flatten()


if __name__ == '__main__':
    try:
        triton_client = httpclient.InferenceServerClient(
                url="127.0.0.1:18000", verbose=False)
    except Exception as e:
        print("channel creation failed: " + str(e))
        sys.exit()
    face_img = cv2.imread('dest.jpg')
    boxes = get_box(triton_client,"dest.jpg")
    print(boxes.shape)
    nms_len =  int(boxes[0]*15+1)
    nms_box = boxes[1:nms_len].reshape(-1,15)
    i = 0
    #print(nms_box.shape)
    for x in nms_box:
        bbox,landmarks = x[0:3],x[5:]
        landmarks = landmarks.reshape(5,2)
        nimg = face_align.norm_crop(face_img, landmarks)
        cv2.imwrite("test_out"+str(i)+'.jpg',nimg)
        i=i+1

    '''
    f = open("demofile2.txt", "a+")
    for box in nms_box:
        f.write('\n')
        for ele in box:
            f.write(str(ele))
            f.write(' ')

    keep = cpu_nms(nms_box,0.6)
    print(keep)
    for i in keep:
        box = nms_box[i]
        if box[4]<0.7:
            continue
        color = (0, 0, 255)
        print(box)
        cv2.rectangle(img, (box[0], box[1]), (box[2], box[3]), color, 2)
        cv2.circle(img, (box[5], box[6]), 1, (0, 0, 255), 4)  # Left eye left corner
        cv2.circle(img, (box[7], box[8]), 1, (0, 255, 255), 4)  # Right eye right corner
        cv2.circle(img, (box[9], box[10]), 1, (255, 0, 255), 4)  # Nose tip
        cv2.circle(img, (box[11], box[12]), 1, (0, 255, 0), 4)  # Left Mouth corner
        cv2.circle(img, (box[13], box[14]), 1, (255, 0, 0), 4)
    cv2.imwrite("img.jpg",img)
    '''

